# -*- coding: utf-8 -*-
import os, sys, cv2, numpy as np, time
import subprocess
from PIL import Image, ImageDraw, ImageFont

from PyQt5.QtCore import QThread, pyqtSignal, QTimer, Qt
from PyQt5.QtGui import QImage, QPixmap
from PyQt5.QtWidgets import QApplication, QMainWindow, QMessageBox

from ui.save_check_ui_V1  import Ui_MainWindow

import resources_rc

# -------------- 全局参数 --------------
DET_MODEL = "onnx_model/yolov5n-face_320_cut.q.onnx"
REC_MODEL = "onnx_model/arcface_mobilefacenet_cut.q.onnx"
FACE_DIR  = "faces"
CAM_ID    = 20
FONT_PATH = "/usr/share/fonts/truetype/noto/NotoSansCJK-Regular.ttc"

os.makedirs(FACE_DIR, exist_ok=True)

TEACHER_DIR = os.path.join(FACE_DIR, "teacher")
STUDENT_DIR = os.path.join(FACE_DIR, "student")
os.makedirs(TEACHER_DIR, exist_ok=True)
os.makedirs(STUDENT_DIR, exist_ok=True)

VERIFY_TEACHER_DIR = TEACHER_DIR
VERIFY_STUDENT_DIR = STUDENT_DIR
VERIFY_THRESHOLD   = 0.55

# -------------- 中文工具 --------------
def put_chinese(img, text, pos, color=(0, 255, 0), size=20):
    try:
        # 保证坐标为整数
        x, y = int(pos[0]), int(pos[1])
        img_pil = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
        draw = ImageDraw.Draw(img_pil)
        font = ImageFont.truetype(FONT_PATH, size, encoding="utf-8")
        draw.text((x, y), text, color[::-1], font)
        return cv2.cvtColor(np.array(img_pil), cv2.COLOR_RGB2BGR)
    except Exception as e:
        # 任何异常直接返回原图
        print("【put_chinese 异常】", e)
        return img

# -------------- 懒加载模型 --------------
def load_detection():
    try:
        from base.detection import Detection
        return Detection(DET_MODEL)
    except Exception as e:
        print("【Detection 加载失败】", e)
        raise

def load_recognition():
    try:
        from base.recognition import Recognition
        return Recognition(REC_MODEL, FACE_DIR)
    except Exception as e:
        print("【Recognition 加载失败】", e)
        raise

# -------------- 后台线程 --------------
class CamThread(QThread):
    frame_ready = pyqtSignal(np.ndarray)
    def __init__(self, cam_id=20):
        super().__init__()
        self.cam_id = cam_id
        self.running = True
    def run(self):
        cap = cv2.VideoCapture(self.cam_id, cv2.CAP_V4L2)
        cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
        cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)
        if not cap.isOpened():
            print("❌ 摄像头无法打开")
            return
        while self.running:
            ok, frame = cap.read()
            if ok:
                self.frame_ready.emit(frame)
            time.sleep(0.01)
        cap.release()
    def stop(self):
        self.running = False
        self.wait()

class DetThread(QThread):
    det_ready = pyqtSignal(np.ndarray, list)
    def __init__(self, det):
        super().__init__()
        self.det = det
        self.running = True
    def feed(self, frame):
        self.latest_frame = frame
    def run(self):
        while self.running:
            if hasattr(self, 'latest_frame'):
                _, boxes = self.det.infer_face(self.latest_frame)
                self.det_ready.emit(self.latest_frame, boxes or [])
            time.sleep(0.01)
    def stop(self):
        self.running = False
        self.wait()

class SaveThread(QThread):
    saved = pyqtSignal(str, str)
    def __init__(self, rec):
        super().__init__()
        self.rec = rec
        self.task = None
        self.running = True
        self.save_root = STUDENT_DIR      # ★ 默认先指向学生
    # ★ 新增接口：动态切换保存目录
    def set_save_root(self, folder):
        self.save_root = folder
    def push_task(self, name, frame, box):
        self.task = (name, frame, box)
    def run(self):
        while self.running:
            if self.task:
                name, frame, box = self.task
                self.task = None
                x1, y1, x2, y2 = map(int, box)
                crop = frame[y1:y2, x1:x2]
                try:
                    vec = self.rec.infer(crop)
                    save_path = os.path.join(self.save_root, f"{name}.npy")
                    np.save(save_path, vec)
                    self.saved.emit(name, save_path)
                except Exception as e:
                    print("【保存失败】", e)
            time.sleep(0.01)
    def stop(self):
        self.running = False
        self.wait()

class VerifyThread(QThread):
    result_ready = pyqtSignal(np.ndarray, dict)

    def __init__(self, det):
        super().__init__()
        self.det = det
        from base.recognition import Recognition
        self.rec = Recognition(REC_MODEL, None)
        self.running = True
        self.latest_frame = None
        self.retry_frame = None

        # ---------- 一次性加载两张人脸库 ----------
        self.face_bank = {}
        for subdir in (VERIFY_TEACHER_DIR, VERIFY_STUDENT_DIR):
            if not os.path.isdir(subdir):
                continue
            for npy in os.listdir(subdir):
                if npy.endswith(".npy"):
                    name = os.path.splitext(npy)[0]
                    try:
                        vec = np.load(os.path.join(subdir, npy))
                        self.face_bank[name] = vec
                    except Exception:
                        pass  # 坏文件直接忽略

        # 如果人脸库为空，线程不会崩，但所有结果都是 unknown
        print(f"[VerifyThread] Loaded {len(self.face_bank)} faces")

    def feed(self, frame):
        self.latest_frame = frame

    def retry(self, frame):
        self.retry_frame = frame

    def run(self):
        while self.running:
            src = self.retry_frame or self.latest_frame
            self.retry_frame = None
            if src is None:
                time.sleep(0.01)
                continue

            try:
                face_imgs, boxes = self.det.infer_face(src)
            except Exception as e:
                print("detection error:", e)
                continue

            results = {}
            if face_imgs is not None and boxes:
                for face_img, box in zip(face_imgs, boxes):
                    vec = None
                    try:
                        vec = self.rec.infer(face_img)
                    except Exception as e:
                        print("recognition error:", e)
                        continue

                    best_name = "unknown"
                    best_score = 0.0
                    for name, v in self.face_bank.items():
                        score = (vec @ v.T)[0][0]
                        if score > best_score and score >= VERIFY_THRESHOLD:
                            best_score, best_name = score, name

                    if best_name != "unknown":
                        prefix = "Teacher" if os.path.isfile(os.path.join(VERIFY_TEACHER_DIR, f"{best_name}.npy")) else "Student"
                        display_name = f"{prefix}-{best_name}"
                    else:
                        display_name = "unknown"
                    results[display_name] = box
            self.result_ready.emit(src, results)
            time.sleep(0.03)

    def stop(self):
        self.running = False
        self.wait()

# -------------- 主窗口 --------------
class MainWindow(QMainWindow, Ui_MainWindow):
    def __init__(self):
        super().__init__()
        self.setupUi(self)
        #窗口美化代码
        self.camera.setStyleSheet("background: transparent;")
        self.tabWidget.setStyleSheet("""
            QTabWidget::pane {
                border-image: url(:/images/bg.jpg) 0 0 0 0 stretch stretch;
            }
        """)
        
        self.pending_save_root = STUDENT_DIR
        self.signed_students = set()
        
        
        
        # 线程先空着
        self.cam_thread  = None
        self.det_thread  = None
        self.save_thread = None
        self.det         = None
        self.rec         = None

        self.latest_frame = None
        self.latest_boxes = []

        # tab2 线程/缓存
        self.verify_cam   = None
        self.verify_det   = None
        self.verify_rec   = None
        self.verify_thr   = None
        self.verify_latest_frame = None
        self.verify_results = {}

        # 定时刷新
        self.timer = QTimer()
        self.timer.timeout.connect(self.update_frame)
        self.timer.start(30)
        
        # ★ 绑定按钮
        self.Ation_input.clicked.connect(self.start_system)
        self.cancle_input.clicked.connect(self.cancel_input)
        self.cancle_input_2.clicked.connect(self.confirm_input)
        # ★ 新增三条
        self.input_teacher.clicked.connect(self.switch_teacher_mode)
        self.input_student.clicked.connect(self.switch_student_mode)
        self.cancle_input_2.clicked.connect(self.open_faces_folder)

        self.pushButton_2.clicked.connect(self.verify_start)
        self.pushButton_3.clicked.connect(self.verify_pause)
        self.pushButton.clicked.connect(self.verify_retry)

    # -------------------------------------------------
    # ★ 新增：切换保存目录
    def switch_teacher_mode(self):
        self.pending_save_root = TEACHER_DIR
        if self.save_thread is not None:
            self.save_thread.set_save_root(TEACHER_DIR)
        self.err_information.append("已切换到【教师】录入模式")

    def switch_student_mode(self):
        self.pending_save_root = STUDENT_DIR
        if self.save_thread is not None:
            self.save_thread.set_save_root(STUDENT_DIR)
        self.err_information.append("已切换到【学生】录入模式")

    def open_faces_folder(self):
        path = os.path.abspath(FACE_DIR)
        if sys.platform == "win32":
            subprocess.Popen(f'explorer "{path}"')
        elif sys.platform == "darwin":
            subprocess.Popen(["open", path])
        else:
            subprocess.Popen(["xdg-open", path])

    def start_system(self):
        if self.cam_thread is not None:
            QMessageBox.information(self, "提示", "系统已启动")
            return
        try:
            self.det = load_detection()
            self.rec = load_recognition()
        except Exception as e:
            QMessageBox.critical(self, "模型加载失败", str(e))
            return

        self.cam_thread  = CamThread(CAM_ID)
        self.det_thread  = DetThread(self.det)
        self.save_thread = SaveThread(self.rec)

        self.cam_thread.frame_ready.connect(self.det_thread.feed)
        self.det_thread.det_ready.connect(self.on_det_ready)
        self.save_thread.saved.connect(self.on_saved)

        self.cam_thread.start()
        self.det_thread.start()
        self.save_thread.start()

        self.err_information.append("摄像头已启动，请正对镜头")

    def on_det_ready(self, frame, boxes):
        self.latest_frame = frame
        self.latest_boxes = boxes
        if boxes and not getattr(self, '_detected_flag', False):
            self._detected_flag = True
            self.err_information.append("已检测到人脸，请您点击【开始录入】！")

    def update_frame(self):
        if self.latest_frame is None:
            return
        vis = self.latest_frame.copy()
        for x1, y1, x2, y2 in self.latest_boxes:
            cv2.rectangle(vis, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 2)
        h, w, ch = vis.shape
        qt_img = QImage(vis.data, w, h, ch * w, QImage.Format_RGB888).rgbSwapped()
        self.camera.setPixmap(QPixmap.fromImage(qt_img))

    def confirm_input(self):
        if self.cam_thread is None:
            QMessageBox.warning(self, "提示", "请先点击【开始录入】启动系统！")
            return
        name = self.input_information.toPlainText().strip()
        if not name:
            QMessageBox.warning(self, "提示", "请输入姓名！")
            return
        if not self.latest_boxes:
            QMessageBox.warning(self, "提示", "未检测到人脸！")
            return
        self.save_thread.push_task(name, self.latest_frame, self.latest_boxes[0])
        self.input_information.append(f"正在保存 {name} ...")

    def cancel_input(self):
        # 1. 清空界面
        self.input_information.clear()
        self.err_information.append("已取消录入")

        # 2. 如果摄像头已启动，就关闭
        if self.cam_thread and self.cam_thread.isRunning():
            self.cam_thread.stop()
        if self.det_thread and self.det_thread.isRunning():
            self.det_thread.stop()
        if self.save_thread and self.save_thread.isRunning():
            self.save_thread.stop()

        # 3. 可选：释放 OpenCV 句柄（CamThread 里已做）
        self.latest_frame = None
        self.latest_boxes = []
        self.camera.clear()          # 让画面黑掉

    def on_saved(self, name, save_path):
        self.input_information.append(f"人脸特征已保存到 {save_path}")

    def closeEvent(self, event):
        if self.cam_thread:  self.cam_thread.stop()
        if self.det_thread:  self.det_thread.stop()
        if self.save_thread: self.save_thread.stop()
        event.accept()

    def verify_start(self):
        if self.verify_thr is not None:
            QMessageBox.information(self, "提示", "验证已启动")
            return
        try:
            self.verify_det = load_detection()
            _ = load_recognition()  # 仅确保模型 OK，不真正使用
        except Exception as e:
            QMessageBox.critical(self, "模型加载失败", str(e))
            return

        self.verify_cam = CamThread(CAM_ID)
        self.verify_thr = VerifyThread(self.verify_det)

        self.verify_cam.frame_ready.connect(self.verify_thr.feed)
        self.verify_thr.result_ready.connect(self.on_verify_result)

        self.verify_cam.start()
        self.verify_thr.start()
        self.textEdit.append("验证摄像头已启动")

    def verify_pause(self):
        if self.verify_cam is None:
            return
        # 1. 停线程
        self.verify_cam.stop()
        self.verify_thr.stop()
        # 2. 清空 QLabel
        self.cmaera2.clear()
        # 3. 立即把画面置黑（可选）
        black = QPixmap(640, 480)  # 尺寸与 QLabel 保持一致
        black.fill(Qt.black)
        self.cmaera2.setPixmap(black)

        self.verify_cam = self.verify_thr = None
        self.textEdit.append("验证已暂停")

    def verify_retry(self):
        if self.verify_latest_frame is None:
            self.textEdit.append("当前无画面，无法重识别")
            return
        self.verify_thr.retry(self.verify_latest_frame)
        self.textEdit.append("已提交重识别任务")

    def on_verify_result(self, frame, results):
        self.verify_latest_frame = frame
        self.verify_results = results

        if frame is None:
            return

        vis = frame.copy()
        for name, box in results.items():
            x1, y1, x2, y2 = map(int, box)
            cv2.rectangle(vis, (x1, y1), (x2, y2), (0, 255, 0), 2)

        # 防止 y 坐标越界
            text_y = max(y1 - 10, 10)
            vis = put_chinese(vis, name, (x1, max(y1 - 35, 5)), size=20, color=(0, 255, 0))

        # 学生签到（仅一次）
            if name.startswith("Student-"):
                short = name.split("-", 1)[1]
                if short not in self.signed_students:
                    self.signed_students.add(short)
                    self.textEdit.append(f"{short} 已签到")

        h, w, ch = vis.shape
        qt_img = QImage(vis.data, w, h, ch * w, QImage.Format_RGB888).rgbSwapped()
        self.cmaera2.setPixmap(QPixmap.fromImage(qt_img))

    # 退出时一并清理
    def closeEvent(self, event):
        if self.verify_cam: self.verify_cam.stop()
        if self.verify_thr: self.verify_thr.stop()
        super().closeEvent(event)

# -------------- 入口 --------------
if __name__ == "__main__":
    try:
        app = QApplication(sys.argv)
        win = MainWindow()
        win.show()
        sys.exit(app.exec_())
    except Exception as e:
        import traceback
        traceback.print_exc()
        input("按回车退出...")
